/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import alloy.ast.ASTDepthFirstReplacer;
import alloy.ast.ArgList;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.Expr;
import alloy.ast.InvocablePara;
import alloy.ast.Invocation;
import alloy.ast.MultiplicityExpr;
import alloy.ast.Node;
import alloy.ast.QuantifiedFormula;
import alloy.ast.SigExpr;
import alloy.ast.SumIntExpr;
import alloy.ast.TypedExpr;
import alloy.ast.VarCreator;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.util.Dbg;


/**
 * A general purpose visitor for replacing variables and in a subtree of the AST.
 * Given a {@link Map} from {@link String} variable names to
 * {@link Expr expressions} that should replace the variables, this visitor does the
 * appropriate replacements in the nodes it encounters.  If a variable name
 * which needs to be replaced is shadowed by an encountered
 * node, the variables with that name are not replaced in the body of that node.
 */
public class ExprAndTypeReplaceVisitor extends ASTDepthFirstReplacer {

    /**
     * mapping from {@link String} names of variables
     * which need to be replaced to
     * {@link Expr}s which should replace them
     */
    private Map _varExprMapping;

    /**
     * characterization of expressions being replaced.
     * @see alloy.ast.Node#annotate
     */
    private String _replaced;

    /**
     * characterization of expressions which are being
     * substituted in
     * @see alloy.ast.Node#annotate
     */
    private String _replacer;

    /** flag indicating whether anything was actually replaced **/
    private boolean _didReplace;

    /**
     * Used to create a new visitor which will then be applied
     * to the body of an invocable.
     * @param invocable_ the invocable which is being invoked
     * @param invocation_ the invocation of the invocable
     * @param params_ parameters passed by the invocation
     */
    public ExprAndTypeReplaceVisitor(InvocablePara invocable_, Invocation invocation_) {
        //Dbg.info(function_.getName().orignodeString());
        _varExprMapping = _getFormalMapping(invocable_, invocation_);
        _replaced = "formal";
        _replacer = "parameter";
	_didReplace = false;
    }

    /**
     * Used to create visitors to visit arbitrary subtrees of
     * the AST.
     */
    public ExprAndTypeReplaceVisitor(Map varExprMapping_,
                                     String replaced_, String replacer_) {
        _varExprMapping = varExprMapping_;
        _replaced = replaced_;
        _replacer = replacer_;
	_didReplace = false;
    }

    /** @return true if at least one expr was actually replaced **/
    public boolean replacementMade() {
	return _didReplace;
    }

    /** visit a var. creator node.  Handle the node specially since
     * it introduces new names and may therefore created shadowing of
     * names of expressions which are being replaced.
     */
    private Object _visit(VarCreator varCreator_) {
        // keep copy of original var. expr mapping
        Map oldVarExprMapping = new HashMap(_varExprMapping);
        Iterator declIter = varCreator_.getDecls().getDeclIter();
        while (declIter.hasNext()) {
            Decl decl = (Decl)declIter.next();
            // visit RHS of decl using current mapping
            decl.setExpr((MultiplicityExpr)decl.getExpr().applyReturnVisitor(this));
            Iterator varIter = decl.getVariables().getVariableIter();
            while (varIter.hasNext()) {
                Variable var = (Variable)varIter.next();
                // see if var name shadows a name of a variable being replaced
                String varName = var.nodeString();
                if (_varExprMapping.containsKey(varName)) {
                    // remove shadowed name's entry from mapping
                    _varExprMapping.remove(varName);
                }
            }
        }
        // visit body of var. creator
        varCreator_.setBody((Node)varCreator_.getBody().applyReturnVisitor(this));
        // restore old var expr mapping
        _varExprMapping = oldVarExprMapping;
        return varCreator_;
    }

    /**
     * @see #_visit(VarCreator)
     */
    public Object visit(QuantifiedFormula quantFormula_) {
        return _visit((VarCreator)quantFormula_);
    }

    /**
     * @see #_visit(VarCreator)
     */
    public Object visit(ComprehensionExpr compExpr_) {
        return _visit((VarCreator)compExpr_);
    }

    /**
     * @see #_visit(VarCreator)
     */
    public Object visit(SumIntExpr sumIntExpr_) {
        return _visit((VarCreator)sumIntExpr_);
    }

    /**
     * Checks if a variable expression should be replaced, and if so
     * replaces it.  Otherwise, modifies the expression's type if
     * necessary.
     * @param varExpr_ the variable expression
     */
    public Object visit(VariableExpr varExpr_) {
        String varString = varExpr_.getVar().getId().nodeString();
        //Dbg.info(varString);
        if (_varExprMapping.containsKey(varString)) {
            Expr replacement = (Expr)((Expr)_varExprMapping.get(varString)).copy();
            replacement.annotateTransformReplacing("replacement of " + _replaced + " " + varExpr_.nodeString()
                                                   + " with " + _replacer + " " + replacement.nodeString(), varExpr_);
	    _didReplace = true;
            return replacement;
        } else {
            return visit((TypedExpr)varExpr_);
        }
    }


    /**
     * Checks if expression needs to be replaced, and if so
     * performs the replacement.  If no replacement is necessary,
     * visits children and sets type of expression properly.
     */
    public Object visit(SigExpr sigExpr_) {
        String exprString = sigExpr_.nodeString();
        //Dbg.info(varString);
        if (_varExprMapping.containsKey(exprString)) {
            Expr replacement = (Expr)((Expr)_varExprMapping.get(exprString)).copy();
            replacement.annotateTransformReplacing("replacement of " + _replaced + " " + sigExpr_.nodeString()
                                                   + " with " + _replacer + " " + replacement.nodeString(), sigExpr_);
	    _didReplace = true;
            return replacement;
        } else {
            return visit((TypedExpr)sigExpr_);
        }
    }

    public Object visit(TypedExpr typedExpr_){
        return visit((Node)typedExpr_);
    }

    /**
     * Get a mapping from formals to parameter
     * @param invocable_ the invocable which is being invoked
     * @param invocation_ the invocation of the invocable
     * @param params_ the parameters passed by the invocation, including implicit
     *         result arg if any
     * @return the mapping
     */
    private Map _getFormalMapping(InvocablePara invocable_, Invocation invocation_) {
        Map formalMapping = new HashMap();

        //Dbg.info("hashcode " + invocation_.hashCode());
        // first, do arguments
        ArgList.InOrderFormalIter declIter = invocable_.getArgList().orderedIter();
        Iterator exprIter = invocation_.getArgs().iterator();
        // add mappings for all arguments
        while (declIter.hasNext()) {
            Variable var = declIter.nextVar();
            //Dbg.info(var.getId().orignodeString());
            Dbg.chk(exprIter.hasNext(), "problem in invocation " + invocation_.nodeString()
                    + " at " + invocation_.getLocation().toString());
            formalMapping.put(var.getId().nodeString(), ((Expr)exprIter.next()).copy());
        }

        Dbg.chk(!exprIter.hasNext(), "wrong number of arguments");
        return formalMapping;
    }
}
